"""
集成：
1、vote: 简单投票
2、average: 平均
"""
import os
import json
import numpy as np


def voting(predict_labels: np.ndarray, num_classes, output_path, id2labels: dict):
    num_samples = predict_labels.shape[1]

    if len(predict_labels.shape) == 3:
        predict_labels = np.argmax(predict_labels, axis=-1)
    assert len(predict_labels.shape) == 2

    res = np.zeros((num_samples, num_classes), dtype=np.int)
    for model in predict_labels:
        for index, label in enumerate(model):
            res[index, label] += 1
    print(res)
    res = np.argmax(res, axis=1).tolist()
    print(res)

    if output_path:
        directory = os.path.dirname(output_path)
        if not os.path.exists(directory):
            os.makedirs(directory)

        output = []
        for i in range(num_samples):
            output.append({"id": i+1, "label": id2labels[res[i]]})
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(output, f)


def average(logits_list, num_classes):
    result = np.zeros((len(logits_list[0]), num_classes))
    for i in logits_list:
        result += i
    result = result / len(logits_list)
    result = np.argmax(result, axis=1)
    return result.tolist()


if __name__ == '__main__':
    # example 1
    # a = np.array([[0.3, 0.2, 0.5], [0.5, 0.1, 0.4], [0.3, 0.6, 0.1]])  # [2, 0, 1]
    # b = np.array([[0.25, 0.25, 0.5], [0.6, 0.2, 0.2], [0.2, 0.7, 0.1]])  # [2, 0, 1]
    # c = np.array([[0.2, 0.5, 0.3], [0.1, 0.1, 0.8], [0.7, 0.2, 0.1]])  # [1, 2, 0]
    # result = voting(np.array([a, b, c]), num_classes=3)
    # print(result)
    # result = average([a,b,c], num_classes=3)
    # print(result)

    # -------------------------------------------------------------------
    # example 2
    # a = np.array([2, 0, 1])
    # b = np.array([2, 0, 1])
    # c = np.array([1, 2, 0])
    # print(voting(np.array([a, b, c]), num_classes=3))
    # -------------------------------------------------------------------

    # example 3
    label_dict_path = r'../data/label_dict.json'
    files_path = [
        # r'../data_split/pseudo_labels/usual_result0.789.txt',
        # r'../data_split/pseudo_labels/usual_result0.7792.txt',
        # r'../data_split/pseudo_labels/usual_result0.7812.txt',
        # r'../data_split/pseudo_labels/usual_result0.7842.txt',
        # r'../data_split/pseudo_labels/usual_result0.7748.txt',
        # r'../data_split/pseudo_labels/virus_result0.6625.txt',
        # r'../data_split/pseudo_labels/virus_result0.6671.txt',
        # r'../data_split/pseudo_labels/virus_result0.6848.txt',
    ]
    with open(label_dict_path, 'r', encoding='utf-8') as f:
        label_dict = json.load(f)
    predicts = []
    for p in files_path:
        with open(p, 'r', encoding='utf-8') as f:
            temp = json.load(f)
        temp = [label_dict[i['label']] for i in temp]
        predicts.append(temp)
    convert_dict = {value: key for key, value in label_dict.items()}
    voting(np.array(predicts, dtype=np.int),
           num_classes=len(label_dict),
           output_path=r'./virus_result.txt',
           id2labels=convert_dict)

